// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/dynamic_library_loader.h>
#include <pollux/expression/simple_function_registry.h>
#include <pollux/functions/function_registry.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

namespace kumo::pollux::functions::test {
    class DynamicLinkTest : public FunctionBaseTest {
    };

    namespace {
        std::string getLibraryPath(const std::string &filename) {
            return fmt::format(
                "{}/{}{}",
                POLLUX_TEST_DYNAMIC_LIBRARY_PATH,
                filename,
                POLLUX_TEST_DYNAMIC_LIBRARY_PATH_SUFFIX);
        }
    } // namespace

    TEST_F(DynamicLinkTest, dynamicLoadFunc) {
        const auto dynamicFunction = [&]() {
            return evaluateOnce<int64_t>("dynamic()", make_row_vector(ROW({}), 1));
        };

        const auto dynamicFunctionNestedCall = [&]() {
            return evaluateOnce<int64_t>(
                "mod(dynamic(), 10)", make_row_vector(ROW({}), 1));
        };

        POLLUX_ASSERT_THROW(
            dynamicFunction(), "Scalar function doesn't exist: dynamic.");

        auto signaturesBefore = get_function_signatures().size();
        std::string libraryPath = getLibraryPath("libpollux_function_dynamic");
        load_dynamic_library(libraryPath);
        auto signaturesAfter = get_function_signatures().size();
        EXPECT_EQ(signaturesAfter, signaturesBefore + 1);
        EXPECT_EQ(123, dynamicFunction());

        auto &registry = exec::simple_functions();
        auto resolved = registry.resolve_function("dynamic", {});
        EXPECT_EQ(TypeKind::BIGINT, resolved->type()->kind());

        EXPECT_EQ(3, dynamicFunctionNestedCall());
    }

    TEST_F(DynamicLinkTest, dynamicLoadSameFuncTwice) {
        const auto dynamicFunction = [&]() {
            return evaluateOnce<int64_t>("dynamic()", make_row_vector(ROW({}), 1));
        };
        auto &registry = exec::simple_functions();
        std::string libraryPath = getLibraryPath("libpollux_function_dynamic");

        load_dynamic_library(libraryPath);
        auto resolved = registry.resolve_function("dynamic", {});
        EXPECT_EQ(TypeKind::BIGINT, resolved->type()->kind());

        auto signaturesBefore = get_function_signatures().size();

        load_dynamic_library(libraryPath);
        auto signaturesAfterSecond = get_function_signatures().size();
        EXPECT_EQ(signaturesAfterSecond, signaturesBefore);
        auto resolvedAfterSecond = registry.resolve_function("dynamic", {});
        EXPECT_EQ(TypeKind::BIGINT, resolvedAfterSecond->type()->kind());
    }

    TEST_F(DynamicLinkTest, dynamicLoadOverwriteFunc) {
        const auto dynamicIntFunction = [&]() {
            return evaluateOnce<int64_t>(
                "dynamic_overwrite()", make_row_vector(ROW({}), 1));
        };
        const auto dynamicVarcharFunction = [&]() {
            return evaluateOnce<std::string>(
                "dynamic_overwrite()", make_row_vector(ROW({}), 1));
        };

        auto &registry = exec::simple_functions();
        auto signaturesBefore = get_function_signatures().size();

        POLLUX_ASSERT_THROW(
            dynamicVarcharFunction(),
            "Scalar function doesn't exist: dynamic_overwrite.");

        std::string libraryPath =
                getLibraryPath("libpollux_overwrite_varchar_function_dynamic");
        load_dynamic_library(libraryPath);
        auto signaturesAfterFirst = get_function_signatures().size();
        EXPECT_EQ(signaturesAfterFirst, signaturesBefore + 1);
        EXPECT_EQ("123", dynamicVarcharFunction());
        auto resolved = registry.resolve_function("dynamic_overwrite", {});
        EXPECT_EQ(TypeKind::VARCHAR, resolved->type()->kind());

        POLLUX_ASSERT_THROW(
            dynamicIntFunction(),
            "Expression evaluation result is not of expected type: dynamic_overwrite() -> CONSTANT vector of type VARCHAR");

        std::string libraryPathInt =
                getLibraryPath("libpollux_overwrite_int_function_dynamic");
        load_dynamic_library(libraryPathInt);

        // The first function loaded should be overwritten.
        POLLUX_ASSERT_THROW(
            dynamicVarcharFunction(),
            "Expression evaluation result is not of expected type: dynamic_overwrite() -> CONSTANT vector of type BIGINT");
        EXPECT_EQ(123, dynamicIntFunction());
        auto signaturesAfterSecond = get_function_signatures().size();
        EXPECT_EQ(signaturesAfterSecond, signaturesAfterFirst);
        auto resolvedAfterSecond = registry.resolve_function("dynamic_overwrite", {});
        EXPECT_EQ(TypeKind::BIGINT, resolvedAfterSecond->type()->kind());
    }

    TEST_F(DynamicLinkTest, dynamicLoadErrFunc) {
        const auto dynamicFunctionFail = [&](const std::optional<int64_t> a,
                                             std::optional<int64_t> b) {
            return evaluateOnce<int64_t>("dynamic_err(c0)", a, b);
        };

        const auto dynamicFunctionSuccess =
                [&](const kumo::pollux::RowVectorPtr &arr) {
            return evaluateOnce<int64_t>("dynamic_err(c0)", arr);
        };

        auto signaturesBefore = get_function_signatures().size();
        POLLUX_ASSERT_THROW(
            dynamicFunctionFail(0, 0), "Scalar function doesn't exist: dynamic_err.");

        std::string libraryPath = getLibraryPath("libpollux_function_err_dynamic");
        load_dynamic_library(libraryPath);

        auto signaturesAfter = get_function_signatures().size();
        EXPECT_EQ(signaturesAfter, signaturesBefore + 1);

        // Expecting a fail because we are not passing in an array.
        POLLUX_ASSERT_THROW(
            dynamicFunctionFail(0, 0),
            "Scalar function signature is not supported: dynamic_err(BIGINT). Supported signatures: (array(bigint)) -> bigint.");

        auto check = make_row_vector(
            {
                make_nullable_array_vector(std::vector<std::vector<std::optional<int64_t> > >{
                    {0, 1, 3, 4, 5, 6, 7, 8, 9}
                })
            });

        // Expecting a success because we are passing in an array.
        EXPECT_EQ(9, dynamicFunctionSuccess(check));
    }

    TEST_F(DynamicLinkTest, dynamicLoadOverloadFunc) {
        const auto dynamicIntFunction = [&](const std::optional<int64_t> a) {
            return evaluateOnce<int64_t>("dynamic_overload(c0)", a);
        };
        const auto dynamicVarcharFunction = [&](const std::optional<std::string> a) {
            return evaluateOnce<std::string>("dynamic_overload(c0)", a);
        };

        auto &registry = exec::simple_functions();
        auto signaturesBefore = get_function_signatures().size();

        POLLUX_ASSERT_THROW(
            dynamicVarcharFunction("1"),
            "Scalar function doesn't exist: dynamic_overload.");

        std::string libraryPath =
                getLibraryPath("libpollux_overload_varchar_function_dynamic");
        load_dynamic_library(libraryPath);
        auto signaturesAfterFirst = get_function_signatures().size();
        EXPECT_EQ(signaturesAfterFirst, signaturesBefore + 1);
        EXPECT_EQ("1", dynamicVarcharFunction("1"));
        auto resolved = registry.resolve_function("dynamic_overload", {VARCHAR()});
        EXPECT_EQ(TypeKind::VARCHAR, resolved->type()->kind());

        // We expect no dynamic_5 with int arguments yet.
        POLLUX_ASSERT_THROW(
            dynamicIntFunction(0),
            "Scalar function signature is not supported: dynamic_overload(BIGINT). Supported signatures: (varchar) -> varchar.");

        std::string libraryPathInt =
                getLibraryPath("libpollux_overload_int_function_dynamic");
        load_dynamic_library(libraryPathInt);

        // The first function loaded should NOT be overwritten.
        EXPECT_EQ("0", dynamicVarcharFunction("0"));
        EXPECT_EQ(0, dynamicIntFunction(0));
        auto signaturesAfterSecond = get_function_signatures().size();
        EXPECT_EQ(signaturesAfterSecond, signaturesAfterFirst);
        auto resolvedAfterSecond =
                registry.resolve_function("dynamic_overload", {BIGINT()});
        EXPECT_EQ(TypeKind::BIGINT, resolvedAfterSecond->type()->kind());
    }
} // namespace kumo::pollux::functions::test
